import fix_data,torch
import numpy as np
from torch import nn, sigmoid

#超参数
lr = 0.0001
input_dim = 40
out_dim = 3
epochs = 600

#搭建神经网络

class csgoAINet(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.layer1 = nn.Sequential(
            nn.Linear(input_dim,20),
            nn.ReLU()
        )
        self.layer2 = nn.Sequential(
            nn.Linear(20,15),
            nn.ReLU(),
            nn.Dropout(0.5)
        )
        self.layer3 = nn.Sequential(
            nn.Linear(15,15),
            nn.ReLU()
        )
        self.layer4 = nn.Sequential(
            nn.Linear(15,out_dim)
        )
    def forward(self,x):
        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.layer4(x)
        return x

model = csgoAINet()
#训练
opt = torch.optim.Adam(model.parameters(),lr)
MSEloss = nn.MSELoss()
#准备数据
data = fix_data.start_file2arrary("data/data-98871.txt")
train_data,val_data = fix_data.start_arrary2train_val(data)
x_train,y_train = fix_data.start_data2x_y(train_data)
x_val,y_val = fix_data.start_data2x_y(val_data)
#转化为numpy并且归一化
x_train = fix_data.minmaxscaler(x_train)
x_val = fix_data.minmaxscaler(x_val)
y_train = np.array(y_train)
y_val = np.array(y_val)
#转化为张量
x_train,y_train,x_val,y_val = torch.from_numpy(x_train).float(),torch.from_numpy(y_train).float(),torch.from_numpy(x_val).float(),torch.from_numpy(y_val).float()
#开始
print("开始训练！")
running_loss = 0.0
running_acc = 0.0
for epoch in range(epochs):
    epoch += 1
    for i,x_data in enumerate(x_train):
        out = model(x_data)
        opt.zero_grad()
        target = y_train[i]
        loss = MSEloss(out,target)
        running_loss = loss.data
        loss.backward()
        opt.step()
    print("目前为第{}训练轮,Loss:{}".format(epoch,running_loss))
#测试
print("开始测试")
model.eval()
running_loss = 0.0
running_acc = 0.0
for epoch in range(epochs):
    epoch += 1
    for i,x_data in enumerate(x_val):
        out = model(x_data)
        opt.zero_grad()
        loss = MSEloss(out,y_val[i])
        running_loss = loss.data
        loss.backward()
        opt.step()
    print("目前为第{}测试轮,Loss:{}".format(epoch,running_loss))
#保存
torch.save(model, "test.pt")

